import datetime
import logging
from typing import List, Dict, Optional
from absl import app, flags
import sys
import collections
import os
import traceback
import time
import subprocess
from concurrent.futures import ProcessPoolExecutor
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt

from coin.base.datetime_util import (to_datetime, to_timestamp_int)
from coin.proto.coin_market_query_pb2 import ProductTicker
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.support.proto_log.logic.util import (read_strat_log, StratInfo)
import python.experimental.lucasyoo.accounting.accounting as lac

sys.path.append(os.path.expandvars('$COIN_DEPLOY_REPO'))
os.environ['keyword'] = 'sfty_delta'  # just for hack. No meaning
import python.coin_deploy.vmm.fabfile as vmm_fabfile

matplotlib.use("Agg")
os.environ["TZ"] = "GMT"
time.tzset()

FLAGS = flags.FLAGS

Data = collections.namedtuple("Data", ["ts", "product_str", "price", "reserve", "hypo_reserve"])

PREV_MACHINE_INFO = {
    "vmm_joy_delta_hedge": "strategy-191.ap-northeast-1",
    "vmm_fncy_delta_hedge": "strategy-191.ap-northeast-1",
    "vmm_grnd_delta_hedge": "strategy-191.ap-northeast-1"
}

ROOT_DIR = '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log'


class OnLog(object):
  def __init__(self,
               start_time: datetime.datetime,
               end_time: datetime.datetime,
               strategy: str,
               ticker_intv_min: float = 1.0):
    self._start_time: datetime.datetime = start_time
    self._end_time: datetime.datetime = end_time
    self._strategy: str = strategy
    self._base = \
      f'{self._strategy.replace("vmm_", "").replace("_delta_hedge", "").split("_")[0].upper()}'
    self._get_hypo_reserve_from_git()
    self._get_base_and_symbol()
    self._kline_list: List[Data] = []
    self._query_interval_feed()
    self._last_reserve: Dict[str, Optional[float]] = {}
    self._last_price: Dict[str, Optional[float]] = {}
    self._last_price_update: Dict[str, Optional[datetime.datetime]] = {}
    self._ticker_intv: datetime.timedelta = datetime.timedelta(minutes=ticker_intv_min)

  def _query_interval_feed(self):
    final_df = pd.DataFrame()
    start_dt = self._start_time
    end_dt = self._end_time
    for product_str in self._product_strs:
      kline_df = lac.query_interval_feed(product_str, start_dt, end_dt)
      final_df = pd.concat([final_df, kline_df])
    if final_df.empty:
      return
    final_df['ts'] = pd.to_datetime(final_df['klineTimestamp'])
    sliced_df = final_df[['ts', 'product_str', 'close']]
    sliced_df.rename(columns={'close': 'price'}, inplace=True)
    sliced_df['reserve'] = np.nan
    sliced_df['hypo_reserve'] = np.nan
    # make sliced_df to list of Data namedtuples
    final_df = sliced_df.apply(lambda x: Data(*x), axis=1).tolist()
    self._kline_list = final_df

  def _get_hypo_reserve_from_git(self):
    self._pb_list = lac.query_coin_deploy_executor_reserve_from_git(self._strategy, self._base)

  def _get_base_and_symbol(self):
    self._product_strs = {data.product_str for data in self._pb_list}
    symbols = {product_str.split(":")[2] for product_str in self._product_strs}
    assert len(symbols) == 1, f'[{self._strategy}] {symbols}'
    self._symbol = symbols.pop()

  def _handle_reserve(self, ts: int, pb: StrategyLog):
    for each_balance in pb.reserve.each_balance:
      if each_balance.currency == self._base:
        total = each_balance.total
        product_str = f"{pb.reserve.market_type}:{pb.reserve.exchange}:{self._symbol}"
        if self._last_reserve.get(product_str) is None or self._last_reserve[product_str] != total:
          logging.info(f'[{product_str}] reserve changed: {to_datetime(ts)}'
                       f' from {self._last_reserve.get(product_str, 0.0)} to {total}')
          self._pb_list.append(Data(to_datetime(ts), product_str, np.nan, total, np.nan))
          self._last_price_update[product_str] = None
        self._last_reserve[product_str] = total

  def _handle_ticker(self, ts: int, pb: StrategyLog):
    for each_ticker in pb.ticker.each_ticker:
      if each_ticker.symbol == self._symbol:
        product_str = f"{pb.ticker.market_type}:{pb.ticker.exchange}:{self._symbol}"
        price = each_ticker.last
        dt = to_datetime(ts)
        if self._last_price.get(product_str) is None or self._last_price_update.get(
            product_str) is None or dt - self._last_price_update[product_str] >= self._ticker_intv:
          self._pb_list.append(Data(dt, product_str, price, np.nan, np.nan))
          self._last_price_update[product_str] = dt
        self._last_price[product_str] = price

  def on_log(self, ts: int, log: bytes):
    pb = StrategyLog()
    pb.ParseFromString(log)
    if pb.type == StrategyLog.RESERVE:
      self._handle_reserve(ts, pb)
    elif pb.type == StrategyLog.TICKER:
      self._handle_ticker(ts, pb)

  def get_corp_offset_pnl(self) -> float:
    final_list = self._kline_list + self._pb_list
    df = pd.DataFrame(final_list).sort_values('ts').reset_index(drop=True)
    df[['price', 'reserve',
        'hypo_reserve']] = df.groupby(['product_str'])[['price', 'reserve',
                                                        'hypo_reserve']].ffill()
    df['reserve'] = df.groupby(['product_str'])[['reserve']].bfill()
    df = df[(df['ts'] >= self._start_time) & (df['ts'] < self._end_time)]
    if FLAGS.strategy_name == "vmm_sfty_delta_hedge":
      df = df[df['ts'] >= datetime.datetime.strptime("2023-04-05 09:00:00", "%Y-%m-%d %H:%M:%S")]
    df['price_diff'] = df.groupby(['product_str'])['price'].diff()
    df['corp_offset_pnl_diff'] = df['price_diff'] * (df['reserve'] - df['hypo_reserve'])
    if self._strategy == 'vmm_wemix_delta_hedge':   # BRL conversion
      df['corp_offset_pnl_diff'] *= 0.2
    df['corp_offset_pnl_diff'].fillna(0.0, inplace=True)
    df['corp_offset_pnl'] = df.groupby(['product_str'])['corp_offset_pnl_diff'].cumsum()
    df['corp_pos'] = (df['reserve'] - df['hypo_reserve'])
    return df

  def plot_corp_pnl_and_offset(self):
    df = self.get_corp_offset_pnl()
    output_dir = os.path.join(os.path.expandvars('$COIN_REPO'), FLAGS.output_dir)
    os.makedirs(output_dir, exist_ok=True)

    df.to_csv(os.path.join(
        output_dir, f'{self._strategy}_offset_pnl_{self._start_time.strftime("%Y%m%d")}'
        f'_{self._end_time.strftime("%Y%m%d")}.csv'),
              index=False)
    # df = pd.read_csv(os.path.join(
    #     output_dir, f'{self._strategy}_offset_pnl_{self._start_time.strftime("%Y%m%d")}'
    #     f'_{self._end_time.strftime("%Y%m%d")}.csv'))
    # df['ts'] = pd.to_datetime(df['ts'])
    # breakpoint()
    daily_df = df.set_index('ts').groupby('product_str').resample('D').last().drop(
        'product_str', axis=1).reset_index().sort_values(['product_str', 'ts'])
    total_corp_offset_pnl = daily_df.groupby(['ts'])['corp_offset_pnl'].sum()
    total_corp_pos = daily_df.groupby(['ts'])['corp_pos'].sum()
    product_strs = daily_df['product_str'].unique()
    plt.rcParams['figure.figsize'] = (30, 20)
    plt.subplot(311)
    plt.plot(total_corp_offset_pnl.index,
             total_corp_offset_pnl.to_numpy(),
             'k-',
             lw=0.5,
             drawstyle='steps-post',
             label='TOTAL')
    title1 = f'corp offset pnl: TOTAL: {total_corp_offset_pnl.iloc[-1]:,.0f}, '
    for product_str in product_strs:
      tmp_df = daily_df[daily_df['product_str'] == product_str]
      plt.plot(tmp_df['ts'],
               tmp_df['corp_offset_pnl'],
               lw=0.5,
               drawstyle='steps-post',
               label=product_str)
      title1 += f"{product_str}: {tmp_df['corp_offset_pnl'].iloc[-1]:,.0f}, "
    plt.legend()
    plt.title(f"{title1} ({total_corp_offset_pnl.index.min()}-{total_corp_offset_pnl.index.max()})")
    plt.tight_layout()
    plt.subplot(312)
    plt.plot(total_corp_pos.index,
             total_corp_pos.to_numpy(),
             'k-',
             lw=0.5,
             drawstyle='steps-post',
             label='TOTAL')

    title2 = f'corp pos: total: {total_corp_pos.iloc[-1]:,.0f}, '
    for product_str in product_strs:
      tmp_df = daily_df[daily_df['product_str'] == product_str]
      plt.plot(tmp_df['ts'], tmp_df['corp_pos'], lw=0.5, drawstyle='steps-post', label=product_str)
      title2 += f"{product_str}: {tmp_df['corp_pos'].iloc[-1]:,.0f}, "
    plt.legend()
    plt.title(f"{title2} ({total_corp_pos.index.min()}-{total_corp_pos.index.max()})")
    plt.tight_layout()
    plt.subplot(313)
    for product_str in product_strs:
      tmp_df = daily_df[daily_df['product_str'] == product_str]
      plt.plot(tmp_df['ts'], tmp_df['price'], lw=0.5, drawstyle='steps-post', label=product_str)
    plt.legend()
    plt.title('price')
    plt.tight_layout()
    plt.savefig(os.path.join(
        output_dir, f'{self._strategy}_offset_pnl_{self._start_time.strftime("%Y%m%d")}'
        f'_{self._end_time.strftime("%Y%m%d")}.png'),
                dpi=300)
    plt.close()


def run(strat_name, machine, start_time, end_time, ticker_intv_min=1.0):
  try:
    onlog = OnLog(start_time, end_time, strat_name, ticker_intv_min=ticker_intv_min)
    if strat_name in PREV_MACHINE_INFO:
      logging.info(
          f'[{strat_name}] reading previous machine logs ({PREV_MACHINE_INFO[strat_name]}->{machine})'
      )
      prev_machine = PREV_MACHINE_INFO[strat_name]
      strat_info_cls = StratInfo(strategy_name=strat_name, machine=prev_machine, trading_date=None)
      read_strat_log(strat_info=strat_info_cls,
                     start_time=start_time,
                     end_time=end_time,
                     callback=onlog.on_log,
                     root_dir=ROOT_DIR)
    strat_info_cls = StratInfo(strategy_name=strat_name, machine=machine, trading_date=None)
    read_strat_log(strat_info=strat_info_cls,
                   start_time=start_time,
                   end_time=end_time,
                   callback=onlog.on_log,
                   root_dir=ROOT_DIR)
    onlog.plot_corp_pnl_and_offset()
  except Exception:
    logging.error(f'[{strat_name}] failed to run', exc_info=True)


def main(_):
  subprocess.run(['git', 'switch', 'master'], cwd=os.path.expandvars('$COIN_DEPLOY_REPO'))
  subprocess.run(['git', 'pull'], cwd=os.path.expandvars('$COIN_DEPLOY_REPO'))
  start_time = datetime.datetime.strptime(FLAGS.start_time, '%Y%m%d')
  end_time = datetime.datetime.strptime(FLAGS.end_time, '%Y%m%d') + datetime.timedelta(days=1)
  strategy_list = vmm_fabfile.strategy_list
  if FLAGS.strategy_name == 'ALL':
    with ProcessPoolExecutor(max_workers=len(strategy_list)) as executor:
      for strat_info in strategy_list:
        strat_name = strat_info['name']
        if strat_name in ["vmm_maicoin_test", "vmm_xrp_sandbox_delta_hedge", "vmm_xrp_sandbox_manipulator", "vmm_grnd_btc_delta_hedge", "vmm_orb_btc_delta_hedge"]:
          continue
        machine = strat_info['host']
        executor.submit(run, strat_name, machine, start_time, end_time)
  else:
    strat_info = [x for x in strategy_list if x['name'] == FLAGS.strategy_name]
    assert len(strat_info) == 1, f'cannot find strategy {FLAGS.strategy_name} in vmm fabfile'
    strat_info = strat_info[0]
    strat_name = strat_info['name']
    machine = strat_info['host']
    run(strat_name, machine, start_time, end_time)


if __name__ == '__main__':
  flags.DEFINE_string('start_time', '20230101', '%Y%m%d')
  flags.DEFINE_string('end_time', '20230331', '%Y%m%d')
  flags.DEFINE_string('strategy_name', 'ALL',
                      'strategy name e.g.vmm_ccx_delta_hedge, if ALL, run every vmm_*_delta_hedge')
  flags.DEFINE_string('output_dir', 'sim_result/accounting', 'output dir')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
